# catgirl_tools.py（整合所有基础功能）
import requests
import json
import os
import logging

# 配置工具类日志
logging.basicConfig(
    level=logging.DEBUG,
    format='[调试] %(asctime)s - %(module)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

# -------------------------- 1. Ollama 基础交互 --------------------------
def check_ollama_alive(url="http://localhost:11434"):
    """检查 Ollama 服务是否启动"""
    try:
        return requests.get(f"{url}/").status_code == 200
    except requests.exceptions.ConnectionError:
        return False

def call_ollama(prompt, model="Catgirl_Assistant", stream=False, url="http://localhost:11434"):
    """调用 Ollama 生成内容"""
    if not check_ollama_alive(url):
        err_msg = "喵~ Ollama 服务没启动哦，先运行 `ollama serve` 再试吧~"
        logging.warning(err_msg)
        return err_msg

    try:
        resp = requests.post(
            f"{url}/api/generate",
            json={"model": model, "prompt": prompt, "stream": stream},
            timeout=30
        )
        if resp.status_code == 200:
            result = resp.json().get("response", "生成内容为空啦~")
            logging.debug(f"Ollama 调用成功，结果长度：{len(result)}")
            return result
        else:
            err_msg = f"Ollama 调用失败，状态码{resp.status_code}~"
            logging.error(err_msg)
            return err_msg
    except Exception as e:
        err_msg = f"调用出错啦：{str(e)}"
        logging.error(err_msg)
        return err_msg

# -------------------------- 2. RPG 配置文件生成 --------------------------
def extract_pure_rpg_config(raw_content):
    """提取纯RPG配置（过滤猫娘闲聊，保留核心格式）"""
    lines = [line.rstrip() for line in raw_content.split('\n') if line.strip()]
    config_lines = []
    rpg_core_fields = {"Id:", "Data:", "Display:", "Lore:"}
    found_fields = set()
    lore_list_active = False
    chat_features = ["“", "”", "希望", "符合你的需求", "觉得", "推荐", "记得", "要不要"]

    for line in lines:
        stripped_line = line.strip()
        current_is_config = False

        if not config_lines:
            if (stripped_line.endswith(":") and len(stripped_line.split(":")) == 2 and stripped_line not in rpg_core_fields) or \
                    (stripped_line.startswith("物品名:") and len(stripped_line.split(":")) >= 2):
                if stripped_line.startswith("物品名:"):
                    item_name = stripped_line.split(":", 1)[1].strip()
                    config_lines.append(f"{item_name}:")
                else:
                    config_lines.append(line)
                current_is_config = True
                logging.debug(f"识别物品名行：{stripped_line}")

        elif any(field in stripped_line for field in rpg_core_fields):
            found_fields.add(next(f for f in rpg_core_fields if f in stripped_line))
            config_lines.append(line)
            current_is_config = True
            lore_list_active = stripped_line == "Lore:"
            logging.debug(f"识别核心字段：{stripped_line}，已找到字段：{found_fields}")

        elif lore_list_active and stripped_line.startswith("- "):
            if not stripped_line.endswith(("喵~", "呀~", "呢~")) or len(stripped_line) > 5:
                config_lines.append(line)
                current_is_config = True
                logging.debug(f"保留Lore项：{stripped_line}")

        if any(f in stripped_line for f in chat_features) and not current_is_config:
            logging.debug(f"过滤闲聊行：{stripped_line}")
            continue

        if found_fields == rpg_core_fields and lore_list_active and not current_is_config:
            lore_count = len([l for l in config_lines if l.strip().startswith("- ")])
            if lore_count < 1:
                logging.debug(f"Lore项不足（{lore_count}项），继续解析")
                continue
            logging.debug(f"配置解析完成，找到所有核心字段，Lore项数量：{lore_count}")
            break

    pure_config = fix_rpg_config_format(config_lines)
    logging.debug(f"提取的纯配置：{pure_config[:300]}...")
    return pure_config

def fix_rpg_config_format(config_lines):
    """修复 RPG 格式缩进"""
    fixed = []
    core_indent = "  "
    list_indent = "    "
    lore_need_indent = False
    rpg_core_fields = {"Id:", "Data:", "Display:", "Lore:"}

    for line in config_lines:
        stripped = line.strip()
        if stripped.endswith(":") and len(stripped.split(":")) == 2 and stripped not in rpg_core_fields:
            fixed.append(stripped)
        elif any(f in stripped for f in ["Id:", "Data:", "Display:", "Lore:"]):
            if ":" in stripped:
                field, val = stripped.split(":", 1)
                fixed.append(f"{core_indent}{field.strip()}: {val.strip()}")
            else:
                fixed.append(f"{core_indent}{stripped}")
            lore_need_indent = stripped == "Lore:"
        elif lore_need_indent and stripped.startswith("- "):
            fixed.append(f"{list_indent}{stripped}")
        else:
            if stripped:
                fixed.append(line)

    return "\n".join([l for l in fixed if l.strip()])

def generate_rpg_file(prompt, output_dir="."):
    """生成 RPG 配置文件并保存"""
    analysis_prompt = f"""
    分析主人需求，判断需要从知识库中获取哪些信息（例如颜色代码、物品属性等）。
    只需返回需要查询的关键词列表（用中文，每行一个，不要解释），如果不需要任何信息就返回空列表。
    
    主人的需求：{prompt}
    """
    logging.debug(f"需求分析提示词：{analysis_prompt[:300]}...")
    needed_knowledge = call_ollama(analysis_prompt).strip()
    logging.debug(f"AI判断需要的知识库信息：{needed_knowledge}")

    knowledge_content = []
    if needed_knowledge and "空列表" not in needed_knowledge:
        keywords = [line.strip() for line in needed_knowledge.splitlines() if line.strip()]
        for keyword in keywords:
            retrieved = ragflow_retrieve(question=keyword)
            if "错误" not in retrieved and "没启动" not in retrieved:
                knowledge_content.append(f"【关于{keyword}的信息】\n{retrieved}")
            else:
                knowledge_content.append(f"【关于{keyword}的信息】检索失败：{retrieved}")

    knowledge_context = "\n\n".join(knowledge_content) if knowledge_content else "无需要参考的知识库信息"
    rpg_prompt = f"""
    生成 RPG 物品配置文件，严格按以下格式：
    【格式】
    物品名:
      Id: 英文大写+下划线
      Data: 0
      Display: '§+颜色代码+物品名' 
      Lore:
        - '内容1'
        - '内容2'
        - '内容3'
    【要求】
    1. 必须优先使用提供的【知识库参考内容】
    2. 配置文件中不允许出现闲聊内容
    3. 确保格式正确，特别是缩进和引号
    
    【知识库参考内容】
    {knowledge_context}
    
    主人需求：{prompt}
    """
    logging.debug(f"RPG生成提示词（含动态知识）：{rpg_prompt[:500]}...")

    raw_content = call_ollama(rpg_prompt)
    logging.debug(f"Ollama返回原始内容：{raw_content[:500]}...")
    if "错误" in raw_content or "没启动" in raw_content:
        return raw_content, ""

    pure_config = extract_pure_rpg_config(raw_content)
    if not pure_config:
        logging.warning("配置提取失败，返回空配置")
        return "生成的配置格式不对哦，再跟我说说需求吧~", raw_content

    item_name = pure_config.split("\n")[0].strip().rstrip(":")
    filename = os.path.join(output_dir, f"{item_name}.yml")
    try:
        with open(filename, "w", encoding="utf-8") as f:
            f.write(pure_config)
        logging.info(f"配置文件已保存至：{filename}")
        return f"配置文件生成好啦~ 保存路径：{filename}", pure_config
    except Exception as e:
        err_msg = f"文件保存失败：{str(e)}"
        logging.error(err_msg)
        return err_msg, pure_config

# -------------------------- 3. RagFlow 知识库检索 --------------------------
def get_ragflow_key():
    """RagFlow API 密钥"""
    return "ragflow-E5MjA3YmI4OTVjOTExZjBiYzNiYTJkZm"  # 请替换为真实密钥

def check_ragflow_alive(url="http://localhost:80"):
    """检查 RagFlow 服务是否启动"""
    try:
        return requests.get(f"{url}/api/v1/version", timeout=5).status_code == 200
    except requests.exceptions.ConnectionError:
        return False

def ragflow_retrieve(question, dataset_ids=["eb7c54a2956911f0be8636bfaf29e47a"], url="http://localhost:80"):
    """检索知识库内容"""
    if not check_ragflow_alive(url):
        err_msg = "RagFlow 服务没启动哦，先启动 Docker 容器再查吧~"
        logging.warning(err_msg)
        return err_msg

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {get_ragflow_key()}"
    }
    payload = {
        "question": question,
        "dataset_ids": dataset_ids,
        "similarity_threshold": 0.2,
        "top_k": 5,
        "highlight": True
    }
    logging.debug(f"知识库检索请求：{question[:100]}...，数据集：{dataset_ids}")

    try:
        resp = requests.post(
            f"{url}/api/v1/retrieval",
            data=json.dumps(payload),
            headers=headers,
            timeout=15
        )
        resp.raise_for_status()
        result = resp.json()

        if result.get("code") != 0:
            err_msg = f"检索失败：{result.get('message', '未知错误')}"
            logging.error(err_msg)
            return err_msg

        chunks = result.get("data", {}).get("chunks", [])
        if not chunks:
            msg = "知识库中没找到相关内容呢~"
            logging.debug(msg)
            return msg

        retrieved = []
        for i, chunk in enumerate(chunks, 1):
            retrieved.append(f"【相关内容{i}】（来源：{chunk.get('document_name', '未知文档')}，相似度：{chunk.get('similarity', 0.0):.2f}）\n{chunk.get('content', '')}")
        logging.debug(f"知识库检索成功，返回{len(chunks)}条结果")
        return "\n\n".join(retrieved)
    except Exception as e:
        err_msg = f"检索出错：{str(e)}"
        logging.error(err_msg)
        return err_msg